import numpy as np
import cv2 as cv

# 读取视频
cap = cv.VideoCapture("image/DOG.wmv")

# 指定追踪目标
ret, frame = cap.read() # 获取第一帧图像
r, h, c, w = 197, 141, 0, 208
win = (c, r, w, h)
roi = frame[r:r+h, c:c+w]

# 计算直方图
hsv_roi = cv.cvtColor(roi, cv.COLOR_BGR2HSV)
roi_hist = cv.calcHist([hsv_roi], [0], None, [180], [0, 180])
# 归一化
cv.normalize(roi_hist, roi_hist, 0, 255, cv.NORM_MINMAX)

"""
Meanshift的API是：
    cv.meanShift(probImage, window, criteria)
参数：
    probImage: ROI区域，即目标的直方图的反向投影
    window： 初始搜索窗口，就是定义ROI的rect
    criteria: 确定窗口搜索停止的准则，主要有迭代次数达到设置的最大值，窗口中心的漂移值大于某个设定的限值等。
    
实现Meanshift的主要流程是：
    1.读取视频文件：cv.videoCapture()
    2.感兴趣区域设置：获取第一帧图像，并设置目标区域，即感兴趣区域
    3.计算直方图：计算感兴趣区域的HSV直方图，并进行归一化
    4.目标追踪：设置窗口搜索停止条件，直方图反向投影，进行目标追踪，并在目标位置绘制矩形框。
"""

# 目标追踪
# 窗口终止条件：最大迭代次数，窗口中心漂移最小值
term = (cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 1)

while(True):
    ret, frame = cap.read()
    if ret == True:
        # 计算直方图的反向投影
        hst = cv.cvtColor(frame, cv.COLOR_BGR2HSV)
        dst = cv.calcBackProject([hst], [0], roi_hist, [0, 180], 1)
        # 进行meanshift最终
        ret, win = cv.meanShift(dst, win, term)
        # 将追踪位置绘制在视频上，并进行显示
        x, y, w, h = win
        img2 = cv.rectangle(frame, (x, y), (x+w, y+h), 255, 2)
        cv.imshow("frame", img2)
        if cv.waitKey(60) & 0xFF == ord("q"):
            break

# 释放资源
cap.release()
cv.destroyAllWindows()